//
//  LlamaServer.swift
//  Sidekick
//
//  Created by Bean John on 10/9/24.
//

import EventSource
import Foundation
import FSKit_macOS
import OSLog
import SimilaritySearchKit

/// The inference server where LLM inference happens
public actor LlamaServer {
    
    // MARK: - Core Configuration
    
    static let logger: Logger = .init(
        subsystem: Bundle.main.bundleIdentifier!,
        category: String(describing: LlamaServer.self)
    )
    
    var modelType: ModelType
    var systemPrompt: String
    var contextLength: Int
    
    // MARK: - Connection Details
    
    var host: String = "127.0.0.1"
    var port: String
    var scheme: String = "http"
    
    // MARK: - Process State
    
    var isStartingServer: Bool = false
    var monitor: Process = Process()
    var process: Process = Process()
    
    /// Tracks request-scoped streaming resources so multiple in-flight calls can run concurrently.
    var activeRequests: [UUID: ActiveRequestContext] = [:]
    var pendingCancellationForAllRequests: Bool = false
    
    // MARK: - Initialization
    
    init(
        modelType: ModelType,
        systemPrompt: String = InferenceSettings.systemPrompt,
        contextLength: Int = InferenceSettings.contextLength
    ) {
        self.modelType = modelType
        self.systemPrompt = systemPrompt
        self.contextLength = contextLength
        self.port = {
            switch modelType {
                case .regular:
                    return "4579"
                case .worker:
                    return "9830"
                case .completions:
                    return "1623"
            }
        }()
    }
    
    // MARK: - Model Metadata
    
    var modelUrl: URL? {
        switch self.modelType {
            case .regular:
                return Settings.modelUrl
            case .worker:
                return InferenceSettings.workerModelUrl
            case .completions:
                return InferenceSettings.completionsModelUrl
        }
    }
    
    var modelName: String {
        return self.modelUrl?.deletingPathExtension().lastPathComponent ?? "Unknown Model"
    }
    
    func modelHasVision(
        usingRemoteModel: Bool
    ) -> Bool {
        return Self.modelHasVision(
            type: self.modelType,
            usingRemoteModel: usingRemoteModel
        )
    }
    
    static func modelHasVision(
        type: ModelType,
        usingRemoteModel: Bool
    ) -> Bool {
        // Return false if using non-regular model
        if type != .regular {
            return false
        }
        // Using local model and has vision enabled, and has projector
        if !usingRemoteModel && InferenceSettings.localModelHasVision {
            return true
        }
        print("\(InferenceSettings.serverModelName) vision: \(InferenceSettings.serverModelHasVision)")
        // Else, get toggle value
        return InferenceSettings.serverModelHasVision
    }
    
    // MARK: - Configuration Mutations
    
    public func setSystemPrompt(_ systemPrompt: String) {
        self.systemPrompt = systemPrompt
    }
    
    /// Function executed when output finishes
    /// - Parameter text: The output generated by the LLM
    public func onFinish(text: String) {}
    
}

extension LlamaServer {
    
    /// Container for request-local streaming state that must be cancelled independently.
    final class ActiveRequestContext {
        
        let id: UUID
        let eventSource: EventSource
        let dataTask: EventSource.DataTask
        let session: URLSession
        var isCancelled: Bool
        
        init(
            id: UUID,
            eventSource: EventSource,
            dataTask: EventSource.DataTask,
            session: URLSession,
            isCancelled: Bool = false
        ) {
            self.id = id
            self.eventSource = eventSource
            self.dataTask = dataTask
            self.session = session
            self.isCancelled = isCancelled
        }
        
        func cancel() {
            guard !isCancelled else { return }
            isCancelled = true
            dataTask.cancel(urlSession: session)
        }
        
    }
    
}
